# -*- coding: utf-8 -*-
# @Time    : 2020/6/20 下午9:27
# @Author  : caotian
# @FileName: lenetmyopiatrain.py
# @Software: PyCharm

import os
import sys
curpath=os.path.abspath(os.curdir)
sys.path.append(curpath)
import lenetmyopiadataloader as lmd
import myopiamodel as mm

import random
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D,Pool2D,Linear
import numpy as np

traindatadir='./data/PALM-Training400/PALM-Training400'
validationdatadir='./data/PALM-Validation400'
csvfile='./data/PALM-Validation-GT/labels.csv'

def train(model):
    with fluid.dygraph.guard():
        print('start training ...')
        model.train()
        epoch_num=5
        opt=fluid.optimizer.Momentum(learning_rate=0.001,momentum=0.9,parameter_list=model.parameters())
        train_loader=lmd.data_loader(traindatadir,batch_size=10,mode='train')
        valid_loader=lmd.valid_data_loader(validationdatadir,csvfile)
        for epoch in range(epoch_num):
            for batch_id,data in enumerate(train_loader()):
                xdata,ydata=data
                img=fluid.dygraph.to_variable(xdata)
                label=fluid.dygraph.to_variable(ydata)
                logits=model(img)
                loss=fluid.layers.sigmoid_cross_entropy_with_logits(logits,label)
                avg_loss=fluid.layers.mean(loss)
                if batch_id % 10 ==0:
                    print("epoch {},batch_id {},loss {}".format(epoch,batch_id,avg_loss.numpy()))
                avg_loss.backward()
                opt.minimize(avg_loss)
                model.clear_gradients()
            model.eval()
            accuracies=[]
            losses=[]
            for batch_id,data in enumerate(valid_loader()):
                xdata,ydata=data
                img=fluid.dygraph.to_variable(xdata)
                label=fluid.dygraph.to_variable(ydata)
                logits=model(img)
                pred=fluid.layers.sigmoid(logits)
                loss=fluid.layers.sigmoid_cross_entropy_with_logits(pred,label)
                pred2=pred *(-1.0)+1.0
                pred=fluid.layers.concat([pred2,pred],axis=1)
                acc=fluid.layers.accuracy(pred,fluid.layers.cast(label,dtype='int64'))
                accuracies.append(acc.numpy())
                losses.append(loss.numpy())
            print("validation accuracy/loss:{}/{}".format(np.mean(accuracies),np.mean(losses)))
            model.train()
        fluid.save_dygraph(model.state_dict(),'myopia-mnist')
        fluid.save_dygraph(opt.state_dict(),'myopia-mnist')
def evaluation(model,params_file_path):
    with fluid.dygraph.guard():
        print("start evaluation...")
        model_state_dict,_=fluid.load_dygraph(params_file_path)
        model.load_dict(model_state_dict)
        model.eval()
        eval_loader=lmd.load_data('eval')
        acc_set=[]
        avg_loss_set=[]
        for batch_id,data in enumerate(eval_loader()):
            xdata,ydata=data
            img=fluid.dygraph.to_variable(xdata)
            label=fluid.dygraph.to_variable(ydata)
            prediction,acc=model(img,label)
            loss=fluid.layers.cross_entropy(input=prediction,label=label)
            avg_loss=fluid.layers.mean(loss)
            acc_set.append(float(acc.numpy()))
            avg_loss_set.append(float(avg_loss.numpy()))
        acc_val_mean=np.array(acc_set).mean()
        avg_loss_val_mean=np.array(avg_loss_set).mean()
        print('loss {},acc {}'.format(avg_loss_val_mean,acc_val_mean))

if __name__ == '__main__':
    with fluid.dygraph.guard():
        model=mm.LeNet(num_classes=1)
    train(model)
    








